Spaces:
Sleeping
Sleeping
File size: 5,108 Bytes
87700f1 ec32911 87700f1 ec32911 307f432 ec32911 94627ea ec32911 87700f1 ec32911 87700f1 655a5f6 87700f1 ec32911 bf039a9 ec32911 ac46e37 dfba72c 8a63318 7b45f67 ec32911 609e92c ec32911 87700f1 ec32911 87700f1 ec32911 655a5f6 ec32911 94627ea 87700f1 ec32911 fab33c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import os
import cv2
import math
import torch
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from glob import glob
from PIL import Image
from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.plots import Annotator, colors
from utils.torch_utils import select_device, smart_inference_mode
from utils.general import check_img_size, Profile, non_max_suppression, scale_boxes
weights = "runs/train/best_striped.pt"
data = "data.yaml"
# Load model
device = select_device('cpu')
model = DetectMultiBackend(weights, device=device, dnn=False, data=data, fp16=False)
false_detection_data = glob(os.path.join("false_detection", '*.jpg'))
false_detection_data = [x.replace('\\', '/') for x in false_detection_data]
def resize_image_pil(image, new_width, new_height):
# Convert to PIL image
img = Image.fromarray(np.array(image))
# Get original size
width, height = img.size
# Calculate scale
width_scale = new_width / width
height_scale = new_height / height
scale = min(width_scale, height_scale)
# Resize
resized = img.resize((int(width*scale), int(height*scale)), Image.NEAREST)
# Crop to exact size
resized = resized.crop((0, 0, new_width, new_height))
return resized
def display_false_detection_data(false_detection_data, number_of_samples):
fig = plt.figure(figsize=(10, 10))
x_count = 5
y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
for i in range(number_of_samples):
plt.subplot(y_count, x_count, i + 1)
img = cv2.imread(false_detection_data[i])
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img)
plt.xticks([])
plt.yticks([])
return fig
def inference(input_img, conf_thres, iou_thres, is_false_detection_images=True, num_false_detection_images=10):
im0 = input_img.copy()
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size((640, 640), s=stride) # check image size
bs = 1
# Run inference
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
with dt[0]:
im = letterbox(input_img, imgsz, stride=stride, auto=True)[0] # padded resize
im = im.transpose((2, 0, 1))[::-1]
im = np.ascontiguousarray(im)
im = torch.from_numpy(im).to(model.device)
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
# Inference
with dt[1]:
pred = model(im, augment=False, visualize=False)
# NMS
with dt[2]:
pred = non_max_suppression(pred, conf_thres, iou_thres, None, False, max_det=10)
# Process predictions
for i, det in enumerate(pred): # per image
seen += 1
annotator = Annotator(im0, line_width=2, example=str(model.names))
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
# Write results
for *xyxy, conf, cls in reversed(det):
c = int(cls) # integer class
label = f'{names[c]} {conf:.2f}'
annotator.box_label(xyxy, label, color=colors(c, True))
if is_false_detection_images:
# Plot the misclassified data
misclassified_images = display_false_detection_data(false_detection_data, number_of_samples=num_false_detection_images)
else:
misclassified_images = None
return im0, misclassified_images
title = "YOLOv9 model to detect shirt/tshirt"
description = "A simple Gradio interface to infer on YOLOv9 model and detect tshirt in image"
examples = [["image_1.jpg", 0.25, 0.45, True, 10],
["image_2.jpg", 0.25, 0.45, True, 10],
["image_3.jpg", 0.25, 0.45, True, 10],
["image_4.jpg", 0.25, 0.45, True, 10],
["image_5.jpg", 0.25, 0.45, True, 10],
["image_6.jpg", 0.25, 0.45, True, 10],
["image_7.jpg", 0.25, 0.45, True, 10],
["image_8.jpg", 0.25, 0.45, True, 10],
["image_9.jpg", 0.25, 0.45, True, 10],
["image_10.jpg", 0.25, 0.45, True, 10]]
demo = gr.Interface(inference,
inputs = [gr.Image(width=320, height=320, label="Input Image"),
gr.Slider(0, 1, 0.25, label="Confidence Threshold"),
gr.Slider(0, 1, 0.45, label="IoU Thresold"),
gr.Checkbox(label="Show False Detection"),
gr.Slider(5, 35, value=10, step=5, label="Number of False Detection")],
outputs= [gr.Image(width=640, height=640, label="Output"),
gr.Plot(label="False Detection")],
title=title,
description=description,
examples=examples)
demo.launch() |