SamDaLamb commited on
Commit
f6679f7
·
verified ·
1 Parent(s): 4a2577e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +22 -4
app.py CHANGED
@@ -32,9 +32,14 @@ model.eval()
32
  # model_path = "./last.pt"
33
  # model = torch.jit.load(model_path, map_location=torch.device("cpu"))
34
  # model.eval()
35
- transform=transforms.Compose([
36
- transforms.ToPILImage(),
37
- transforms.Resize((512,640)),
 
 
 
 
 
38
  transforms.ToTensor()
39
  ])
40
  # transform = transforms.Compose([
@@ -45,9 +50,22 @@ transform=transforms.Compose([
45
  OBJECT_NAMES = ['enemies']
46
 
47
  def detect_objects_in_image(image):
 
 
 
 
 
 
48
 
 
 
 
 
 
 
49
  img_tensor = transform(image).unsqueeze(0)
50
- orig_w, orig_h = image.size
 
51
 
52
  with torch.no_grad():
53
  pred = model(img_tensor)[0]
 
32
  # model_path = "./last.pt"
33
  # model = torch.jit.load(model_path, map_location=torch.device("cpu"))
34
  # model.eval()
35
+ # transform=transforms.Compose([
36
+ # transforms.ToPILImage(),
37
+ # transforms.Resize((512,640)),
38
+ # transforms.ToTensor()
39
+ # ])
40
+ transform = transforms.Compose([
41
+ transforms.ToPILImage(), # Ensure input is a PIL image
42
+ transforms.Resize((512, 640)),
43
  transforms.ToTensor()
44
  ])
45
  # transform = transforms.Compose([
 
50
  OBJECT_NAMES = ['enemies']
51
 
52
  def detect_objects_in_image(image):
53
+ """
54
+ Detect objects in the given image.
55
+ """
56
+ # Ensure image is a PIL Image
57
+ if isinstance(image, torch.Tensor):
58
+ image = transforms.ToPILImage()(image) # Convert tensor to PIL image
59
 
60
+ if isinstance(image, Image.Image):
61
+ orig_w, orig_h = image.size # PIL image size returns (width, height)
62
+ else:
63
+ raise TypeError(f"Expected a PIL Image but got {type(image)}")
64
+
65
+ # Apply transformation
66
  img_tensor = transform(image).unsqueeze(0)
67
+
68
+
69
 
70
  with torch.no_grad():
71
  pred = model(img_tensor)[0]