Update app.py

app.py

@@ -14,76 +14,13 @@ torch.hub.download_url_to_file('https://tochkanews.ru/wp-content/uploads/2020/09
 torch.hub.download_url_to_file('https://s.rdrom.ru/1/pubs/4/35893/1906770.jpg', '2.jpg')
 torch.hub.download_url_to_file('https://static.mk.ru/upload/entities/2022/04/17/07/articles/detailPicture/5b/39/28/b6/ffb1aa636dd62c30e6ff670f84474f75.jpg', '3.jpg')
 
-def yolox_inference(
-    image_path: gr.inputs.Image = None,
-    model_path: gr.inputs.Dropdown = 'kadirnar/yolox_s-v0.1.1',
-    config_path: gr.inputs.Textbox = 'configs.yolox_s',
-    image_size: gr.inputs.Slider = 640
-):
-    """
-    YOLOX inference function
-    Args:
-        image: Input image
-        model_path: Path to the model
-        config_path: Path to the config file
-        image_size: Image size
-    Returns:
-        Rendered image
-    """
-
-    model = YoloxDetector(model_path, config_path=config_path, device="cpu", hf_model=True)
-    pred = model.predict(image_path=image_path, image_size=image_size)
-    return pred
-        
-
-inputs = [
-    gr.inputs.Image(type="filepath", label="Input Image"),
-    gr.inputs.Dropdown(
-        label="Model Path",
-        choices=[
-            "kadirnar/yolox_s-v0.1.1",
-            "kadirnar/yolox_m-v0.1.1",
-            "kadirnar/yolox_tiny-v0.1.1",
-        ],
-        default="kadirnar/yolox_s-v0.1.1",
-    ),
-    gr.inputs.Dropdown(
-        label="Config Path",
-        choices=[
-            "configs.yolox_s",
-            "configs.yolox_m",
-            "configs.yolox_tiny",
-        ],
-        default="configs.yolox_s",
-    ),
-    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
-]
-
-outputs = gr.outputs.Image(type="filepath", label="Output Image")
-title = "YOLOX is a high-performance anchor-free YOLO."
-
-examples = [
-    ["1.jpg", "kadirnar/yolox_m-v0.1.1", "configs.yolox_m", 640],
-    ["2.jpg", "kadirnar/yolox_s-v0.1.1", "configs.yolox_s", 640],
-    ["3.jpg", "kadirnar/yolox_tiny-v0.1.1", "configs.yolox_tiny", 640],
-]
-
-demo_app = gr.Interface(
-    fn=yolox_inference,
-    inputs=inputs,
-    outputs=outputs,
-    title=title,
-    examples=examples,
-    cache_examples=True,
-    theme='huggingface',
-)
-demo_app.launch(debug=True, enable_queue=True)
-
 
 COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
           [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
 
 
+
+
 def get_class_list_from_input(classes_string: str):
     if classes_string == "":
         return []
@@ -91,29 +28,6 @@ def get_class_list_from_input(classes_string: str):
     classes_list = [x.strip() for x in classes_list]
     return classes_list
 
-def infer(img, model_name: str, prob_threshold: int, classes_to_show = str):
-    feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
-    model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
-
-    img = Image.fromarray(img)
-    
-    pixel_values = feature_extractor(img, return_tensors="pt").pixel_values
-
-    with torch.no_grad():
-        outputs = model(pixel_values, output_attentions=True)
-    
-    probas = outputs.logits.softmax(-1)[0, :, :-1]
-    keep = probas.max(-1).values > prob_threshold
-
-    target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0)
-    postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
-    bboxes_scaled = postprocessed_outputs[0]['boxes']
-    
-    classes_list = get_class_list_from_input(classes_to_show)
-    res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model, classes_list)
-    
-    return res_img
-
 def plot_results(pil_img, prob, boxes, model, classes_list):
     plt.figure(figsize=(16,10))
     plt.imshow(pil_img)
@@ -142,17 +56,81 @@ def fig2img(fig):
     img = Image.open(buf)
     return img
 
-image_in = gr.components.Image()
-image_out = gr.components.Image()
-model_choice = gr.components.Dropdown(["yolos-tiny", "yolos-small", "yolos-base", "yolos-small-300", "yolos-small-dwr"], value="yolos-small", label="YOLOS Model")
-prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold")
-classes_to_show = gr.components.Textbox(placeholder="e.g. car, truck, traffic light", label="Classes to use (empty means all classes)")
 
-Iface = gr.Interface(
-    fn=infer,
-    inputs=[image_in,model_choice, prob_threshold_slider, classes_to_show],
-    outputs=image_out,
+    
+def inference(
+    image_path: gr.inputs.Image = None,
+    model_path: gr.inputs.Dropdown = 'kadirnar/yolox_s-v0.1.1',
+    image_size: gr.inputs.Slider = 640
+):
+
+    if model_name in ("yolox_s-v0.1.1", "yolox_m-v0.1.1", "yolox_tiny-v0.1.1"): 
+        model = YoloxDetector(f"kadirnar/{model_name}", device="cpu", hf_model=True)
+        pred = model.predict(image_path=image_path, image_size=image_size)
+        return pred
+
+    else: 
+        feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
+        model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
+    
+        img = Image.fromarray(img)
+        
+        pixel_values = feature_extractor(img, return_tensors="pt").pixel_values
+    
+        with torch.no_grad():
+            outputs = model(pixel_values, output_attentions=True)
+        
+        probas = outputs.logits.softmax(-1)[0, :, :-1]
+        keep = probas.max(-1).values > prob_threshold
+    
+        target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0)
+        postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
+        bboxes_scaled = postprocessed_outputs[0]['boxes']
+        
+        classes_list = get_class_list_from_input(classes_to_show)
+        res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model, classes_list)
+        
+        return res_img
+
+        
+
+inputs = [
+    gr.inputs.Image(type="filepath", label="Input Image"),
+    image_in,model_choice, prob_threshold_slider, classes_to_show
+    gr.inputs.Dropdown(
+        label="Model Path",
+        choices=[
+            "yolox_s-v0.1.1",
+            "yolox_m-v0.1.1",
+            "yolox_tiny-v0.1.1",
+            "yolos-tiny",
+            "yolos-small",
+            "yolos-base",
+            "yolos-small-300",
+            "yolos-small-dwr"
+        ],
+        default="kadirnar/yolox_s-v0.1.1",
+    ),
+    gr.inputs.Slider(minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
+    gr.inputs.Textbox(placeholder="e.g. car, truck, traffic light", label="Classes to use (empty means all classes)"),
+]
+
+outputs = gr.outputs.Image(type="filepath", label="Output Image")
+
+examples = [
+    ["1.jpg", "kadirnar/yolox_m-v0.1.1", 0.8, 640, ""],
+    ["2.jpg", "kadirnar/yolox_s-v0.1.1", 0.8, 640, ""],
+    ["3.jpg", "kadirnar/yolox_tiny-v0.1.1", 0.8, 640, ""],
+]
+
+demo_app = gr.Interface(
+    fn=inference,
+    inputs=inputs,
+    outputs=outputs,
+    title="Object Detection with YOLO",
     examples=examples,
-    title="Object Detection with YOLOS",
-    description=description,
-).launch()
+    cache_examples=True,
+    theme='huggingface',
+)
+demo_app.launch(debug=True, enable_queue=True)