Skip to content

Commit

Permalink
Fix formatting
Browse files Browse the repository at this point in the history 
Adjust formatting to new versions of the YOLOv5 scripts. See also: ultralytics#12592
  • Loading branch information
@maxsitt
maxsitt committed Jan 28, 2024
1 parent eb0954b commit cf6711b
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 101 deletions.
146 changes: 73 additions & 73 deletions classify/predict.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,27 +100,27 @@

@smart_inference_mode()
def run(
weights=ROOT / "yolov5s-cls.pt", # model.pt path(s)
source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam)
data=ROOT / "data/coco128.yaml", # dataset.yaml path
imgsz=(224, 224), # inference size (height, width)
device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
nosave=False, # do not save images/videos
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / "runs/predict-cls", # save results to project/name
name="exp", # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
sort_top1=False, # sort images to folders with predicted top1 class as folder name
sort_prob=False, # sort images first by probability and then by top1 class
concat_csv=False, # concatenate metadata .csv files and append classification results
new_csv=False # create new .csv file with classification results
weights=ROOT / "yolov5s-cls.pt", # model.pt path(s)
source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam)
data=ROOT / "data/coco128.yaml", # dataset.yaml path
imgsz=(224, 224), # inference size (height, width)
device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
nosave=False, # do not save images/videos
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / "runs/predict-cls", # save results to project/name
name="exp", # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
sort_top1=False, # sort images to folders with predicted top1 class as folder name
sort_prob=False, # sort images first by probability and then by top1 class
concat_csv=False, # concatenate metadata .csv files and append classification results
new_csv=False # create new .csv file with classification results
):
source = str(source)
save_img = not nosave and not source.endswith(".txt") # save inference images
Expand Down Expand Up @@ -217,13 +217,13 @@ def run(
f.write(text + "\n")

# Write image filename and prediction results to lists
lst_img.append(f'{p.name}')
lst_top1.append(f'{names[top5i[0]]}')
lst_top2.append(f'{names[top5i[1]]}')
lst_top3.append(f'{names[top5i[2]]}')
lst_top1_prob.append(f'{prob[top5i[0]]:.2f}')
lst_top2_prob.append(f'{prob[top5i[1]]:.2f}')
lst_top3_prob.append(f'{prob[top5i[2]]:.2f}')
lst_img.append(f"{p.name}")
lst_top1.append(f"{names[top5i[0]]}")
lst_top2.append(f"{names[top5i[1]]}")
lst_top3.append(f"{names[top5i[2]]}")
lst_top1_prob.append(f"{prob[top5i[0]]:.2f}")
lst_top2_prob.append(f"{prob[top5i[1]]:.2f}")
lst_top3_prob.append(f"{prob[top5i[2]]:.2f}")

# Stream results
im0 = annotator.result()
Expand Down Expand Up @@ -282,75 +282,75 @@ def run(

# Print estimated inference time per image
inference_runtime = time.monotonic() - start_inference
LOGGER.info(f'\nEstimated inference time per image: {round((inference_runtime / len(lst_img)) * 1000, 2)} ms')
LOGGER.info(f"\nEstimated inference time per image: {round((inference_runtime / len(lst_img)) * 1000, 2)} ms")

# Create folder to save results
Path(f'{save_dir}/results').mkdir(parents=True, exist_ok=True)
Path(f"{save_dir}/results").mkdir(parents=True, exist_ok=True)

# Write prediction results to .csv
df_results = pd.DataFrame(
{'img_name': lst_img,
'top1': lst_top1,
'top1_prob': lst_top1_prob,
'top2': lst_top2,
'top2_prob': lst_top2_prob,
'top3': lst_top3,
'top3_prob': lst_top3_prob
{"img_name": lst_img,
"top1": lst_top1,
"top1_prob": lst_top1_prob,
"top2": lst_top2,
"top2_prob": lst_top2_prob,
"top3": lst_top3,
"top3_prob": lst_top3_prob
})
df_results['top1_prob'] = pd.to_numeric(df_results['top1_prob'])
df_results['top2_prob'] = pd.to_numeric(df_results['top2_prob'])
df_results['top3_prob'] = pd.to_numeric(df_results['top3_prob'])
df_results.to_csv(f'{save_dir}/results/pred_results.csv', index=False)
df_results["top1_prob"] = pd.to_numeric(df_results["top1_prob"])
df_results["top2_prob"] = pd.to_numeric(df_results["top2_prob"])
df_results["top3_prob"] = pd.to_numeric(df_results["top3_prob"])
df_results.to_csv(f"{save_dir}/results/pred_results.csv", index=False)

# Write mean classification probability per top 1 class to .csv
df_top1_prob = pd.DataFrame(
{'top1': (df_results['top1'].sort_values()
{"top1": (df_results["top1"].sort_values()
.unique()),
'top1_prob_mean': (df_results.groupby(['top1'])['top1_prob']
"top1_prob_mean": (df_results.groupby(["top1"])["top1_prob"]
.mean()
.round(2)
.reset_index(drop=True))
})
df_top1_prob.to_csv(f'{save_dir}/results/top1_prob_mean.csv', index=False)
df_top1_prob.to_csv(f"{save_dir}/results/top1_prob_mean.csv", index=False)

# Plot boxplot with the classification probability per top 1 class
(df_results.plot(kind='box',
column='top1_prob',
by='top1',
(df_results.plot(kind="box",
column="top1_prob",
by="top1",
rot=90,
ylim=(0, 1),
yticks=([x/10 for x in range(0, 11)]),
figsize=(15, 10),
xlabel='Top 1 class',
ylabel='Classification probability'))
plt.rcParams['axes.axisbelow'] = True
plt.grid(axis='y', color='gray', linewidth=0.5, alpha=0.2)
plt.suptitle('')
plt.title('Classification probability per top 1 class')
plt.savefig(f'{save_dir}/results/top1_prob.png', dpi=300, bbox_inches='tight')
xlabel="Top 1 class",
ylabel="Classification probability"))
plt.rcParams["axes.axisbelow"] = True
plt.grid(axis="y", color="gray", linewidth=0.5, alpha=0.2)
plt.suptitle("")
plt.title("Classification probability per top 1 class")
plt.savefig(f"{save_dir}/results/top1_prob.png", dpi=300, bbox_inches="tight")
plt.close()

# Plot barplot with the mean classification probability per top 1 class
(df_top1_prob.sort_values(by='top1_prob_mean', ascending=False)
.plot(kind='bar',
x='top1',
y='top1_prob_mean',
edgecolor='black',
(df_top1_prob.sort_values(by="top1_prob_mean", ascending=False)
.plot(kind="bar",
x="top1",
y="top1_prob_mean",
edgecolor="black",
rot=90,
ylim=(0, 1),
yticks=([x/10 for x in range(0, 11)]),
figsize=(15, 10),
legend=False,
xlabel='Top 1 class',
ylabel='Mean classification probability',
title='Mean classification probability per top 1 class'))
plt.grid(axis='y', color='gray', linewidth=0.5, alpha=0.2)
plt.savefig(f'{save_dir}/results/top1_prob_mean.png', dpi=300, bbox_inches='tight')
xlabel="Top 1 class",
ylabel="Mean classification probability",
title="Mean classification probability per top 1 class"))
plt.grid(axis="y", color="gray", linewidth=0.5, alpha=0.2)
plt.savefig(f"{save_dir}/results/top1_prob_mean.png", dpi=300, bbox_inches="tight")
plt.close()

# Concatenate all metadata .csv files and add new columns with classification results
if concat_csv:
meta_csv_files = list(Path(source).parent.glob('**/metadata*.csv'))
meta_csv_files = list(Path(source).parent.glob("**/metadata*.csv"))
if len(meta_csv_files) > 0:
df_all = []
for csv in meta_csv_files:
Expand All @@ -359,23 +359,23 @@ def run(
if not df_csv.empty:
df_all.append(df_csv)
except pd.errors.EmptyDataError:
print(f'\nBroken metadata*.csv file with no columns: {csv}')
print(f"\nBroken metadata*.csv file with no columns: {csv}")
df_concat = pd.concat(df_all, ignore_index=True)
df_concat = pd.concat([df_concat, df_results.drop(columns=['img_name'])], axis=1)
df_concat.to_csv(f'{save_dir}/results/{name}_metadata_classified.csv', index=False)
df_concat = pd.concat([df_concat, df_results.drop(columns=["img_name"])], axis=1)
df_concat.to_csv(f"{save_dir}/results/{name}_metadata_classified.csv", index=False)
else:
print('\nCould not find any metadata*.csv files!')
print("\nCould not find any metadata*.csv files!")

# Write classification results to new .csv file and extract timestamp + tracking ID
if new_csv:
df_new = df_results
df_new.insert(1, 'timestamp', df_new['img_name'].str[:24])
df_new.insert(2, 'track_ID', df_new['img_name'].str[10:].str.extract('_(.*)_crop'))
df_new.to_csv(f'{save_dir}/results/{name}_data_classified.csv', index=False)
df_new.insert(1, "timestamp", df_new["img_name"].str[:24])
df_new.insert(2, "track_ID", df_new["img_name"].str[10:].str.extract("_(.*)_crop"))
df_new.to_csv(f"{save_dir}/results/{name}_data_classified.csv", index=False)

# Print script run time
script_runtime = time.monotonic() - start_time
LOGGER.info(f'\nScript run time: {round(script_runtime / 60, 3)} min\n')
LOGGER.info(f"\nScript run time: {round(script_runtime / 60, 3)} min\n")


def parse_opt():
Expand Down
53 changes: 25 additions & 28 deletions classify/val.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def run(
dataloader=None,
criterion=None,
pbar=None,
task='val', # val or test dataset split for validation
task="val", # val or test dataset split for validation
):
# Initialize/load model and set device
training = model is not None
Expand Down Expand Up @@ -117,15 +117,12 @@ def run(
data = Path(data)
task = task if task in ("val", "test") else "val"
if task == "val":
test_dir = data / "val" if (data / "val").exists() else data / "test" # data/val or data/test
test_dir = data / "val" if (data / "val").exists() else data / "test" # data/val or data/test
if task == "test":
test_dir = data / "test" if (data / "test").exists() else data / "val" # data/test or data/val
dataloader = create_classification_dataloader(path=test_dir,
imgsz=imgsz,
batch_size=batch_size,
augment=False,
rank=-1,
workers=workers)
test_dir = data / "test" if (data / "test").exists() else data / "val" # data/test or data/val
dataloader = create_classification_dataloader(
path=test_dir, imgsz=imgsz, batch_size=batch_size, augment=False, rank=-1, workers=workers
)

model.eval()
pred, targets, loss, dt = [], [], 0, (Profile(device=device), Profile(device=device), Profile(device=device))
Expand Down Expand Up @@ -184,12 +181,12 @@ def run(
# Write results to .csv
class_names = list(dict(model.names.items()).values())
df_results = pd.DataFrame(
{'class': ['all'] + class_names,
'images': lst_images,
'top1_acc': lst_top1_acc,
'top5_acc': lst_top5_acc
{"class": ["all"] + class_names,
"images": lst_images,
"top1_acc": lst_top1_acc,
"top5_acc": lst_top5_acc
})
df_results.to_csv(f'{save_dir}/valid_results_{task}.csv', index=False)
df_results.to_csv(f"{save_dir}/valid_results_{task}.csv", index=False)

# Write true classes and predicted classes to list
classes_true = targets.tolist()
Expand All @@ -198,16 +195,16 @@ def run(
# Write metrics to .csv
report = classification_report(classes_true, classes_pred, target_names=class_names, output_dict=True)
df_metrics = (pd.DataFrame(report).transpose()
.drop(['accuracy'])
.rename({'macro avg': 'all', 'weighted avg': 'all_weighted'})
.rename(columns={'support': 'images'})
.astype({'images': int})
.drop(["accuracy"])
.rename({"macro avg": "all", "weighted avg": "all_weighted"})
.rename(columns={"support": "images"})
.astype({"images": int})
.round(3))
df_metrics_all = df_metrics.loc[['all', 'all_weighted']]
df_metrics = (pd.concat([df_metrics_all, df_metrics.drop(['all', 'all_weighted'])])
.reset_index(names='class'))
df_metrics = df_metrics[['class', 'images', 'precision', 'recall', 'f1-score']]
df_metrics.to_csv(f'{save_dir}/valid_metrics_{task}.csv', index=False)
df_metrics_all = df_metrics.loc[["all", "all_weighted"]]
df_metrics = (pd.concat([df_metrics_all, df_metrics.drop(["all", "all_weighted"])])
.reset_index(names="class"))
df_metrics = df_metrics[["class", "images", "precision", "recall", "f1-score"]]
df_metrics.to_csv(f"{save_dir}/valid_metrics_{task}.csv", index=False)

# Plot results as confusion matrix
number_classes = len(class_names)
Expand All @@ -223,14 +220,14 @@ def run(
font_size, font_size_values = 12, 10
else:
font_size, font_size_values = 12, 12
cf_matrix = np.around(confusion_matrix(classes_true, classes_pred, normalize='true'), 3)
cf_matrix = np.around(confusion_matrix(classes_true, classes_pred, normalize="true"), 3)
cf_matrix_plot = ConfusionMatrixDisplay(cf_matrix, display_labels=class_names)
plt.rcParams.update({'font.size': font_size})
cf_matrix_plot.plot(cmap='Blues', xticks_rotation='vertical', values_format='.2g')
plt.rcParams.update({"font.size": font_size})
cf_matrix_plot.plot(cmap="Blues", xticks_rotation="vertical", values_format=".2g")
for values in cf_matrix_plot.text_.ravel():
values.set_fontsize(font_size_values)
cf_matrix_plot.ax_.set_title('Normalized confusion matrix')
plt.savefig(f'{save_dir}/confusion_matrix_{task}.png', dpi=600, bbox_inches='tight')
cf_matrix_plot.ax_.set_title("Normalized confusion matrix")
plt.savefig(f"{save_dir}/confusion_matrix_{task}.png", dpi=600, bbox_inches="tight")
plt.close()

return top1, top5, loss
Expand Down

0 comments on commit cf6711b

Please sign in to comment.