Matharrr commited on
Commit
3cd2468
·
1 Parent(s): d7a1ba5

first commit

Browse files
Files changed (1) hide show
  1. app.py +7 -4
app.py CHANGED
@@ -46,14 +46,16 @@ class CUTModel:
46
  self.model.netG.load_state_dict(state_dict) # Load state_dict langsung ke netG
47
  self.model.eval()
48
 
49
- # Transformasi input
50
  self.transform = transforms.Compose([
51
- transforms.Resize((256, 256)),
52
  transforms.ToTensor(),
53
  transforms.Normalize((0.5,), (0.5,))
54
  ])
55
 
56
  def predict(self, image):
 
 
 
57
  # Preprocess
58
  image_tensor = self.transform(image).unsqueeze(0).to(self.device)
59
 
@@ -61,9 +63,10 @@ class CUTModel:
61
  with torch.no_grad():
62
  output = self.model.netG(image_tensor)
63
 
64
- # Postprocess
65
  output_image = (output.squeeze().cpu().clamp(-1, 1) + 1) / 2.0
66
- return transforms.ToPILImage()(output_image)
 
67
 
68
  # Inisialisasi model
69
  print("Inisialisasi model...")
 
46
  self.model.netG.load_state_dict(state_dict) # Load state_dict langsung ke netG
47
  self.model.eval()
48
 
49
+ # Transformasi input - Tanpa resize
50
  self.transform = transforms.Compose([
 
51
  transforms.ToTensor(),
52
  transforms.Normalize((0.5,), (0.5,))
53
  ])
54
 
55
  def predict(self, image):
56
+ # Simpan ukuran asli
57
+ original_size = image.size
58
+
59
  # Preprocess
60
  image_tensor = self.transform(image).unsqueeze(0).to(self.device)
61
 
 
63
  with torch.no_grad():
64
  output = self.model.netG(image_tensor)
65
 
66
+ # Postprocess - Resize ke ukuran asli
67
  output_image = (output.squeeze().cpu().clamp(-1, 1) + 1) / 2.0
68
+ output_image = transforms.ToPILImage()(output_image).resize(original_size)
69
+ return output_image
70
 
71
  # Inisialisasi model
72
  print("Inisialisasi model...")