Spaces:
Running
Running
first commit
Browse files
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 |
-
|
|
|
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...")
|