Update app.py (#3)

- Update app.py (57093c2dd4ced9ccd5b11c1c3fbee526011cae60)

app.py

@@ -13,10 +13,22 @@ import torch.backends.cudnn as cudnn
 from numpy import random
 import numpy as np
 
+# Add safe globals for PyTorch 2.6+ loading
+import numpy
+torch.serialization.add_safe_globals([
+    numpy.core.multiarray._reconstruct,
+    numpy.core.multiarray.scalar,
+    numpy.core.numeric.True_,
+    numpy.core.numeric._frombuffer,
+    numpy.dtype,
+    numpy._globals._NoValue,
+    numpy.core.numeric.asarray,
+    numpy.ndarray
+])
+
 BASE_DIR = "/home/user/app"
 os.chdir(BASE_DIR)
-os.makedirs(f"{BASE_DIR}/input",exist_ok=True)
-# os.system(f"git clone https://github.com/WongKinYiu/yolov7.git {BASE_DIR}/yolov7")
+os.makedirs(f"{BASE_DIR}/input", exist_ok=True)
 sys.path.append(f'{BASE_DIR}/yolov7')
 os.system("pip install yolov7-package==0.0.12")
 
@@ -35,13 +47,6 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=3):
 
 
 def detect(opt, save_img=False):    
-#     from models.experimental import attempt_load
-#     from utils.datasets import LoadStreams, LoadImages
-#     from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
-#         scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
-#     from utils.plots import plot_one_box
-#     from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
-
     from yolov7_package import Yolov7Detector
     from yolov7_package.models.experimental import attempt_load
     from yolov7_package.utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
@@ -49,6 +54,15 @@ def detect(opt, save_img=False):
     from yolov7_package.utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
     from yolov7_package.utils.datasets import LoadStreams, LoadImages
     
+    # Monkey patch the attempt_load function to use weights_only=False
+    original_torch_load = torch.load
+    def patched_torch_load(w, map_location=None):
+        return original_torch_load(w, map_location=map_location, weights_only=False)
+    
+    # Apply the monkey patch
+    import yolov7_package.models.experimental
+    yolov7_package.models.experimental.torch.load = patched_torch_load
+    
     bbox = {}
     source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace
     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
@@ -65,8 +79,16 @@ def detect(opt, save_img=False):
     half = device.type != 'cpu'  # half precision only supported on CUDA
 
     # Load model
-    det = Yolov7Detector(weights=weights, traced=False)
-    model = attempt_load(weights, map_location=device)  # load FP32 model
+    try:
+        det = Yolov7Detector(weights=weights, traced=False)
+        model = attempt_load(weights, map_location=device)  # load FP32 model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        # Fallback: try loading with weights_only=False directly
+        print("Trying fallback method to load model...")
+        model = torch.load(weights, map_location=device, weights_only=False)
+        det = None
+    
     stride = int(model.stride.max())  # model stride
     imgsz = check_img_size(imgsz, s=stride)  # check img_size
 
@@ -145,7 +167,6 @@ def detect(opt, save_img=False):
             if len(det):
                 # Rescale boxes from img_size to im0 size
                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
-#                 print(f"BOXES ---->>>> {det[:, :4]}")
                 bbox[f"{txt_path.split('/')[4]}"]=(det[:, :4]).numpy()
 
                # Print results
@@ -168,38 +189,17 @@ def detect(opt, save_img=False):
             # Print time (inference + NMS)
             print(f'{s}Done. ({(1E3 * (t2 - t1)):.1f}ms) Inference, ({(1E3 * (t3 - t2)):.1f}ms) NMS')
 
-            # Stream results
-            # if view_img:
-            #     cv2.imshow(str(p), im0)
-            #     cv2.waitKey(1)  # 1 millisecond
-
             # Save results (image with detections)
             if save_img:
                 if dataset.mode == 'image':
-                    # Image.fromarray(im0).show()
                     cv2.imwrite(save_path, im0) 
                     print(f" The image with the result is saved in: {save_path}")
-                # else:  # 'video' or 'stream'
-                #     if vid_path != save_path:  # new video
-                #         vid_path = save_path
-                #         if isinstance(vid_writer, cv2.VideoWriter):
-                #             vid_writer.release()  # release previous video writer
-                #         if vid_cap:  # video
-                #             fps = vid_cap.get(cv2.CAP_PROP_FPS)
-                #             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-                #             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                #         else:  # stream
-                #             fps, w, h = 30, im0.shape[1], im0.shape[0]
-                #             save_path += '.mp4'
-                #         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-                #     vid_writer.write(im0)
 
     if save_txt or save_img:
         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
-        #print(f"Results saved to {save_dir}{s}")
 
     print(f'Done. ({time.time() - t0:.3f}s)')
-    return bbox,save_path
+    return bbox, save_path
 
 class options:
     def __init__(self, weights, source, img_size=640, conf_thres=0.1, iou_thres=0.45, device='', 
@@ -226,19 +226,29 @@ class options:
         self.no_trace=no_trace
 
 def get_output(input_image):
-    ### Numpy -> PIL
-    input_image = Image.fromarray(input_image).convert('RGB')
-    input_image.save(f"{BASE_DIR}/input/image.jpg")
-    source = f"{BASE_DIR}/input"
-    opt = options(weights='logo_detection.pt',source=source)
-    bbox = None
-    with torch.no_grad():
-        bbox,output_path = detect(opt)
-    if os.path.exists(output_path):
-        return Image.open(output_path)
-    else:
+    try:
+        # Numpy -> PIL
+        input_image = Image.fromarray(input_image).convert('RGB')
+        input_image.save(f"{BASE_DIR}/input/image.jpg")
+        source = f"{BASE_DIR}/input"
+        opt = options(weights='logo_detection.pt', source=source)
+        bbox = None
+        with torch.no_grad():
+            bbox, output_path = detect(opt)
+        if os.path.exists(output_path):
+            return Image.open(output_path)
+        else:
+            return input_image
+    except Exception as e:
+        print(f"Error in get_output: {e}")
+        # Return the original image if there's an error
         return input_image
 
 
+# Add debug prints to check environment
+print(f"PyTorch version: {torch.__version__}")
+print(f"Current working directory: {os.getcwd()}")
+print(f"Files in directory: {os.listdir('.')}")
+
 demo = gr.Interface(fn=get_output, inputs="image", outputs="image")
 demo.launch(debug=True)