Spaces:
Runtime error
Runtime error
File size: 2,707 Bytes
965ddf7 dda86de 965ddf7 cf72695 965ddf7 6e130fa 965ddf7 6e130fa 965ddf7 8136a78 965ddf7 cf72695 965ddf7 c8a1260 965ddf7 0f019c5 965ddf7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import streamlit as st
import torch
import torchvision
import cv2
import numpy as np
import torch.nn as nn
from torchvision.ops import box_iou
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
# apply nms algorithm
def apply_nms(orig_prediction, iou_thresh=0.3):
# torchvision returns the indices of the bboxes to keep
keep = torchvision.ops.nms(orig_prediction['boxes'], orig_prediction['scores'], iou_thresh)
final_prediction = orig_prediction
final_prediction['boxes'] = final_prediction['boxes'][keep]
final_prediction['scores'] = final_prediction['scores'][keep]
final_prediction['labels'] = final_prediction['labels'][keep]
return final_prediction
# Draw the bounding box
def plot_img_bbox(img, target):
h,w,c = img.shape
for box in (target['boxes']):
xmin, ymin, xmax, ymax = int((box[0].cpu()/1024)*w), int((box[1].cpu()/1024)*h), int((box[2].cpu()/1024)*w),int((box[3].cpu()/1024)*h)
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
label = "palm"
# Add the label and confidence score
label = f'{label}'
cv2.putText(img, label, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
# Display the image with detections
filename = 'pred.jpg'
cv2.imwrite(filename, img)
# transform image
test_transforms = A.Compose([
A.Resize(height=1024, width=1024, always_apply=True),
A.Normalize(always_apply=True),
ToTensorV2(always_apply=True),])
# select device (whether GPU or CPU)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# model loading
model = torch.load('pickel.pth',map_location=torch.device('cpu'))
model = model.to(device)
st.title("🌴Palm trees detection🌴")
file_name = st.file_uploader("Upload oil palm tree image")
if file_name is not None:
col1, col2 = st.columns(2)
image = np.array(Image.open(file_name))
col1.image(image, use_column_width=True)
transformed = test_transforms(image= image)
image_transformed = transformed["image"]
image_transformed = image_transformed.unsqueeze(0)
image_transformed = image_transformed.to(device)
# inference
model.eval()
with torch.no_grad():
predictions = model(image_transformed)[0]
nms_prediction = apply_nms(predictions, iou_thresh=0.1)
plot_img_bbox(image, nms_prediction)
pred = np.array(Image.open("pred.jpg"))
col2.image(pred, use_column_width=True)
word = "Number of palm trees detected : "+str(len(nms_prediction["boxes"]))
st.write(word)
|