Spaces:
Running
Running
| import gradio as gr | |
| import albumentations as A | |
| from functions import * | |
| warnings.filterwarnings('ignore') | |
| # 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) | |
| #-> Tuple[Dict, float] | |
| def predict(img) : | |
| # Start a timer | |
| start_time = timer() | |
| image = np.array(img) | |
| h,w,_ = image.shape | |
| hw = h*w | |
| if hw < 2*1024*1024: | |
| # Transform the target image and add a batch dimension | |
| #image_transformed = test_transforms() | |
| 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) | |
| pred = plot_img_bbox(image, nms_prediction) | |
| #pred = np.array(Image.open("pred.jpg")) | |
| word = "Number of palm trees detected : "+str(len(nms_prediction["boxes"])) | |
| # Calculate the prediction time | |
| pred_time = round(timer() - start_time, 5) | |
| # Return the prediction dictionary and prediction time | |
| return pred,word | |
| else: | |
| crop(image) | |
| locations = np.load("locations.npy") | |
| n = inference(image,locations,model,test_transforms,device) | |
| # | |
| empty_image = np.zeros(image.shape) | |
| del image | |
| gc.collect() | |
| sleep(1) | |
| word = "Number of palm trees detected : "+str(n) | |
| pred = create_new_ortho(locations,empty_image) | |
| # remove files and folders | |
| os.remove("locations.npy") | |
| shutil.rmtree("images", ignore_errors=True) | |
| shutil.rmtree("labels", ignore_errors=True) | |
| return pred,word | |
| image = gr.components.Image() | |
| out_im = gr.components.Image() | |
| out_lab = gr.components.Label() | |
| ### 4. Gradio app ### | |
| # Create title, description and article strings | |
| title = "🌴Palm trees detection🌴" | |
| description = "Faster r-cnn model to detect oil palm trees in drones images." | |
| article = "Created by data354." | |
| # Create examples list from "examples/" directory | |
| example_list = [["examples/" + example] for example in os.listdir("examples")] | |
| #[gr.Label(label="Predictions"), # what are the outputs? | |
| #gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs | |
| # Create examples list from "examples/" directory | |
| # Create the Gradio demo | |
| demo = gr.Interface(fn=predict, # mapping function from input to output | |
| inputs= image, #gr.Image(type="pil"), # what are the inputs? | |
| outputs=[out_im,out_lab], | |
| examples=example_list, | |
| title=title, | |
| description=description, | |
| article=article | |
| ) | |
| # Launch the demo! | |
| demo.launch(debug = True) |