{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nfrom tqdm.auto import tqdm\nimport shutil as sh\nimport torch\nimport sys\nsys.path.insert(0, '/kaggle/input/yolov5')\nfrom models.experimental import *\nfrom utils.datasets import *","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!pip install --no-deps '/kaggle/input/weighted-boxes-fusion/' > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from ensemble_boxes import *\ndef run_wbf(boxes, scores, image_size=1023, iou_thr=0.5, skip_box_thr=0.7, weights=None):\n    #boxes = [prediction[image_index]['boxes'].data.cpu().numpy()/(image_size-1) for prediction in predictions]\n    #scores = [prediction[image_index]['scores'].data.cpu().numpy() for prediction in predictions]\n    labels = [np.zeros(score.shape[0]) for score in scores]\n    boxes = [box/(image_size) for box in boxes]\n    boxes, scores, labels = weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n    #boxes, scores, labels = nms(boxes, scores, labels, weights=[1,1,1,1,1], iou_thr=0.5)\n    boxes = boxes*(image_size)\n    return boxes, scores, labels\n\ndef TTAImage(image, index):\n    image1 = image.copy()\n    if index==0: \n        rotated_image = cv2.rotate(image1, cv2.ROTATE_90_CLOCKWISE)\n        return rotated_image\n    elif index==1:\n        rotated_image2 = cv2.rotate(image1, cv2.ROTATE_90_CLOCKWISE)\n        rotated_image2 = cv2.rotate(rotated_image2, cv2.ROTATE_90_CLOCKWISE)\n        return rotated_image2\n    elif index==2:\n        rotated_image3 = cv2.rotate(image1, cv2.ROTATE_90_CLOCKWISE)\n        rotated_image3 = cv2.rotate(rotated_image3, cv2.ROTATE_90_CLOCKWISE)\n        rotated_image3 = cv2.rotate(rotated_image3, cv2.ROTATE_90_CLOCKWISE)\n        return rotated_image3\n    elif index == 3:\n        return image1\n    \ndef rotBoxes90(boxes, im_w, im_h):\n    ret_boxes =[]\n    for box in boxes:\n        x1, y1, x2, y2 = box\n        x1, y1, x2, y2 = x1-im_w//2, im_h//2 - y1, x2-im_w//2, im_h//2 - y2\n        x1, y1, x2, y2 = y1, -x1, y2, -x2\n        x1, y1, x2, y2 = int(x1+im_w//2), int(im_h//2 - y1), int(x2+im_w//2), int(im_h//2 - y2)\n        x1a, y1a, x2a, y2a = min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)\n        ret_boxes.append([x1a, y1a, x2a, y2a])\n    return np.array(ret_boxes)\n\ndef detect1Image(im0, imgsz, model, device, conf_thres, iou_thres):\n    img = letterbox(im0, new_shape=imgsz)[0]\n    # Convert\n    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416\n    img = np.ascontiguousarray(img)\n\n\n    img = torch.from_numpy(img).to(device)\n    img =  img.float()  # uint8 to fp16/32\n    img /= 255.0   \n    if img.ndimension() == 3:\n        img = img.unsqueeze(0)\n\n    # Inference\n    pred = model(img, augment=False)[0]\n\n    # Apply NMS\n    pred = non_max_suppression(pred, conf_thres, iou_thres)\n\n    boxes = []\n    scores = []\n    for i, det in enumerate(pred):  # detections per image\n        # save_path = 'draw/' + image_id + '.jpg'\n        if det is not None and len(det):\n            # Rescale boxes from img_size to im0 size\n            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()\n\n            # Write results\n            for *xyxy, conf, cls in det:\n                boxes.append([int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])])\n                scores.append(conf)\n\n    return np.array(boxes), np.array(scores) ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def detect():\n    source = '/kaggle/input/global-wheat-detection/test/'\n    weights = '/kaggle/input/yolov5-pth/last2.0.pt'\n#     if not os.path.exists(weights):\n#         weights = '../input/yolov5/bestv4.pt'\n    imgsz = 1024\n    conf_thres = 0.5\n    iou_thres = 0.6\n    is_TTA = True\n    \n    imagenames =  os.listdir(source)\n    \n    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n    # Load model\n    model = torch.load(weights, map_location=device)['model'].float()  # load to FP32\n    model.to(device).eval()\n    \n    dataset = LoadImages(source, img_size=imgsz)\n\n    results = []\n    fig, ax = plt.subplots(5, 2, figsize=(30, 70))\n    count = 0\n    # img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img\n    #for path, img, im0s, _ in dataset:\n    for name in imagenames:\n        image_id = name.split('.')[0]\n        im01 = cv2.imread('%s/%s.jpg'%(source,image_id))  # BGR\n        assert im01 is not None, 'Image Not Found '\n        # Padded resize\n        im_w, im_h = im01.shape[:2]\n        if is_TTA:\n            enboxes = []\n            enscores = []\n            for i in range(4):\n                im0 = TTAImage(im01, i)\n                boxes, scores = detect1Image(im0, imgsz, model, device, conf_thres, iou_thres)\n                for _ in range(3-i):\n                    boxes = rotBoxes90(boxes, im_w, im_h)\n                    \n                if 1: #i<3:\n                    enboxes.append(boxes)\n                    enscores.append(scores) \n            boxes, scores = detect1Image(im01, imgsz, model, device, conf_thres, iou_thres)\n            enboxes.append(boxes)\n            enscores.append(scores)\n\n            boxes, scores, labels = run_wbf(enboxes, enscores, image_size = im_w, iou_thr=0.6, skip_box_thr=0.5)\n            boxes = boxes.astype(np.int32).clip(min=0, max=im_w)\n        else:\n            boxes, scores = detect1Image(im01, imgsz, model, device, conf_thres, iou_thres)\n\n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        boxes = boxes[scores >= 0.05].astype(np.int32)\n        scores = scores[scores >=float(0.05)]\n        if count<10:\n            #sample = image.permute(1,2,0).cpu().numpy()\n            for box, score in zip(boxes,scores):\n                cv2.rectangle(im0,\n                              (box[0], box[1]),\n                              (box[2]+box[0], box[3]+box[1]),\n                              (220, 0, 0), 2)\n                cv2.putText(im0, '%.2f'%(score), (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX ,  \n                   0.5, (255,255,255), 2, cv2.LINE_AA)\n            ax[count%5][count//5].imshow(im0)\n            count+=1\n            \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n\n        results.append(result)\n    return results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = detect()\ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}