Update app.py

app.py

@@ -56,6 +56,11 @@ def detect_Custom(img,model):
         model = TracedModel(model, device, opt.img_size)
     if half:
         model.half() 
+        
+    classify = False
+    if classify:
+        modelc = load_classifier(name='resnet101', n=2)  # initialize
+        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
     vid_path, vid_writer = None, None
     if webcam:
         view_img = check_imshow()
@@ -82,6 +87,11 @@ def detect_Custom(img,model):
 
         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
         t2 = time_synchronized()
+        
+        
+        # Apply Classifier
+        if classify:
+            pred = apply_classifier(pred, modelc, img, im0s)
     
         for i, det in enumerate(pred): 
             if webcam: 
@@ -160,4 +170,4 @@ examples1=[["Image1.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["Image2.jpeg", "Y
 
 Top_Title="<center>Yolov7 Custom Trained by <a href='https://www.linkedin.com/in/owaiskhan9654/' style='text-decoration: underline' target='_blank'>Owais Ahmad</center></a>"
 
-gr.Interface(detect_Custom,[gr.Image(type="pil"),gr.Dropdown(default="Yolo_v7_Custom_trained_By_Owais",choices=["Yolo_v7_Custom_trained_By_Owais","yolov7","yolov7-e6"])],gr.Image(type="pil"),title=Top_Title,examples=examples1,description=Custom_description,article=Footer,cache_examples=False).launch()
+gr.Interface(detect_Custom,[gr.Image(type="pil"),gr.Dropdown(default="Yolo_v7_Custom_trained_By_Owais",choices=["Yolo_v7_Custom_trained_By_Owais","yolov7","yolov7-e6"])],gr.Image(type="pil"),title=Top_Title,examples=examples1,description=Custom_description,article=Footer).launch()