arnaucas commited on
Commit
13ef97c
·
1 Parent(s): 1282b23

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -2,6 +2,8 @@ import gradio as gr
2
  import os
3
  from transformers import pipeline
4
  from pathlib import Path
 
 
5
 
6
  example_imgs = ["examples/img0.jpg",
7
  "examples/img1.jpg",
@@ -11,8 +13,7 @@ example_imgs = ["examples/img0.jpg",
11
  pipe = pipeline("image-classification", model="arnaucas/wildfire-classifier")
12
 
13
  def inference(image):
14
- if image in example_imgs:
15
- image = "./examples/" + image
16
  output = pipe(image)
17
  result = {item['label']: item['score'] for item in output}
18
  return result
 
2
  import os
3
  from transformers import pipeline
4
  from pathlib import Path
5
+ from PIL import Image
6
+ import numpy as np
7
 
8
  example_imgs = ["examples/img0.jpg",
9
  "examples/img1.jpg",
 
13
  pipe = pipeline("image-classification", model="arnaucas/wildfire-classifier")
14
 
15
  def inference(image):
16
+ image = Image.fromarray(np.uint8(image)).convert('RGB')
 
17
  output = pipe(image)
18
  result = {item['label']: item['score'] for item in output}
19
  return result