Spaces:
Sleeping
Sleeping
| from fastai.vision.all import * | |
| from io import BytesIO | |
| import requests | |
| import streamlit as st | |
| import numpy as np | |
| import torch | |
| import time | |
| import cv2 | |
| from numpy import random | |
| from streamlit_image_select import image_select | |
| from models.experimental import attempt_load | |
| 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 | |
| def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32): | |
| # Resize and pad image while meeting stride-multiple constraints | |
| shape = img.shape[:2] # current shape [height, width] | |
| if isinstance(new_shape, int): | |
| new_shape = (new_shape, new_shape) | |
| # Scale ratio (new / old) | |
| r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) | |
| if not scaleup: # only scale down, do not scale up (for better test mAP) | |
| r = min(r, 1.0) | |
| # Compute padding | |
| ratio = r, r # width, height ratios | |
| new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) | |
| dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding | |
| if auto: # minimum rectangle | |
| dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding | |
| elif scaleFill: # stretch | |
| dw, dh = 0.0, 0.0 | |
| new_unpad = (new_shape[1], new_shape[0]) | |
| ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios | |
| dw /= 2 # divide padding into 2 sides | |
| dh /= 2 | |
| if shape[::-1] != new_unpad: # resize | |
| img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) | |
| top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) | |
| left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) | |
| img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border | |
| return img, ratio, (dw, dh) | |
| def detect_modify(img0, model, conf=0.4, imgsz=640, conf_thres = 0.25, iou_thres=0.45): | |
| st.image(img0, caption="Your image", use_column_width=True) | |
| stride = int(model.stride.max()) # model stride | |
| imgsz = check_img_size(imgsz, s=stride) # check img_size | |
| # Padded resize | |
| img0 = cv2.cvtColor(np.asarray(img0), cv2.COLOR_RGB2BGR) | |
| img = letterbox(img0, imgsz, stride=stride)[0] | |
| # Convert | |
| img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 | |
| img = np.ascontiguousarray(img) | |
| # Get names and colors | |
| names = model.module.names if hasattr(model, 'module') else model.names | |
| colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] | |
| # Run inference | |
| old_img_w = old_img_h = imgsz | |
| old_img_b = 1 | |
| t0 = time.time() | |
| img = torch.from_numpy(img).to(device) | |
| # img /= 255.0 # 0 - 255 to 0.0 - 1.0 | |
| img = img/255.0 | |
| if img.ndimension() == 3: | |
| img = img.unsqueeze(0) | |
| # Inference | |
| # t1 = time_synchronized() | |
| with torch.no_grad(): # Calculating gradients would cause a GPU memory leak | |
| pred = model(img)[0] | |
| # t2 = time_synchronized() | |
| # Apply NMS | |
| pred = non_max_suppression(pred, conf_thres, iou_thres) | |
| # t3 = time_synchronized() | |
| # Process detections | |
| # for i, det in enumerate(pred): # detections per image | |
| gn = torch.tensor(img0.shape)[[1, 0, 1, 0]] # normalization gain whwh | |
| det = pred[0] | |
| if len(det): | |
| # Rescale boxes from img_size to im0 size | |
| det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round() | |
| # Print results | |
| s = '' | |
| for c in det[:, -1].unique(): | |
| n = (det[:, -1] == c).sum() # detections per class | |
| s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string | |
| # Write results | |
| for *xyxy, conf, cls in reversed(det): | |
| label = f'{names[int(cls)]} {conf:.2f}' | |
| plot_one_box(xyxy, img0, label=label, color=colors[int(cls)], line_thickness=1) | |
| f""" | |
| ### Prediction result: | |
| """ | |
| img0 = cv2.cvtColor(np.asarray(img0), cv2.COLOR_BGR2RGB) | |
| st.image(img0, caption="Prediction Result", use_column_width=True) | |
| #set paramters | |
| imgsz = 640 | |
| conf = 0.4 | |
| iou_thres=0.45 | |
| device = torch.device("cpu") | |
| path = "./" | |
| """ | |
| # YOLOv7/ YOLOv7-X | |
| This is a object detection model for Chair, (Lamp, Rest,) Sofa, and Table. | |
| """ | |
| weight_path = './' + st.selectbox('Select Model', | |
| ['yolov7_best', 'yolov7x_best', 'yolov7_finder_best', 'yolov7x_finder_best']) + '.pt' | |
| conf_thres = (st.slider("Confidence Threshold (%)", 0, 100, 40))/100 | |
| # Load model | |
| model = attempt_load(weight_path, map_location=torch.device('cpu')) # load FP32 model | |
| option = st.radio("Select one way to demo: ", ["upload image", "image URL", "or try some preset images"]) | |
| if option == "upload image": | |
| uploaded_file = st.file_uploader("Please upload an image.") | |
| if uploaded_file is not None: | |
| img = PILImage.create(uploaded_file) | |
| detect_modify(img, model, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres) | |
| elif option == "image URL": | |
| url = st.text_input("Please input a url.") | |
| if url != "": | |
| try: | |
| response = requests.get(url) | |
| pil_img = PILImage.create(BytesIO(response.content)) | |
| detect_modify(pil_img, model, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres) | |
| except: | |
| st.text("Problem reading image from", url) | |
| elif option == "or try some preset images": | |
| img_select = image_select( | |
| label="Select a picture to detect", | |
| images=[ | |
| # Chair | |
| "https://www.ikea.com/au/en/images/products/nordviken-chair-antique-stain__0832454_pe777681_s5.jpg", | |
| # Sofa | |
| "https://assets.boconcept.com/b1c0b22e-ef01-4d5b-af4d-ad43018a1f5b/1560164_PNG-Web%2072dpi.png?format=pjpg&auto=webp&fit=bounds&width=3020&quality=75%2C60&height=2265", | |
| # Table | |
| "https://habitt.com/cdn/shop/files/2_2_51e1b37c-8035-4abd-8e93-331f145525f5.jpg?v=1697278017", | |
| # Table | |
| "https://m.media-amazon.com/images/I/51zvHEqiKOL._AC_UF1000,1000_QL80_.jpg", | |
| "https://wpmedia.roomsketcher.com/content/uploads/2021/12/09085551/Living_room_idea_wood_details.jpg", | |
| "https://goodhomes.wwmindia.com/content/2022/jan/living-room-picture-by-studio-noughts.jpg", | |
| "https://media.houseandgarden.co.uk/photos/618946a9eea7137eaf372dee/master/w_1600%2Cc_limit/038-2.jpg", | |
| "https://www.checkatrade.com/blog/wp-content/uploads/2023/10/Feature-navy-living-room.jpg" | |
| ], | |
| captions=["Picture 1", "Picture 2", "Picture 3", "Picture 4", | |
| "Picture 5", "Picture 6", "Picture 7", "Picture 8"],) | |
| if (img_select): | |
| url = str(img_select)[:100] | |
| try: | |
| response = requests.get(url) | |
| pil_img = PILImage.create(BytesIO(response.content)) | |
| detect_modify(pil_img, model, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres) | |
| except: | |
| st.text("Problem reading image from", url) |