Jasur05 commited on
Commit
1f88afb
·
verified ·
1 Parent(s): 6aae550

correct rgb match

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -16,13 +16,15 @@ model = load_model()
16
  transform = transforms.Compose([
17
  transforms.Resize((224, 224)),
18
  transforms.ToTensor(),
19
- transforms.Normalize(mean=[0.5], std=[0.5])
20
  ])
21
 
22
  # Prediction function
23
  def predict(image):
 
 
24
  image = image.convert("RGB")
25
- input_tensor = transform(image).unsqueeze(0) # Add batch dim
26
 
27
  with torch.no_grad():
28
  outputs = model(input_tensor)
 
16
  transform = transforms.Compose([
17
  transforms.Resize((224, 224)),
18
  transforms.ToTensor(),
19
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
20
  ])
21
 
22
  # Prediction function
23
  def predict(image):
24
+ if image is None:
25
+ return {"Error": 1.0}
26
  image = image.convert("RGB")
27
+ input_tensor = transform(image).unsqueeze(0)
28
 
29
  with torch.no_grad():
30
  outputs = model(input_tensor)