Spaces:
Build error
Build error
File size: 2,315 Bytes
1865436 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import pytorch_lightning as pl
import numpy as np
from PIL import Image
from src.ss.datasets_signboard_detection.dataset import Labelizer
class Color_convert():
def __init__(self):
super().__init__()
self.labels = {'bien': "red"}
def transform(self, label):
return self.labels[label]
def num_classes(self):
return len(self.labels)
def compose(output, mask):
w,h = mask.shape
for i in range(0, w):
for j in range(0,h):
if (mask[i,j] > 0.5):
output[i,j] = 1
return output
class POIDetectionTask(pl.LightningModule):
def __init__(self,
model,
score):
super().__init__()
self.model = model
self.output = []
self.score = score
self.labelizer = Labelizer()
self.color_convert = Color_convert()
def forward(self, x):
output = self.model(x)
return output
def predict_step(self, test_batch, batch_idx):
images, targets = test_batch
outputs = self(images)
for target in outputs:
shape = target['boxes']
masks = target['masks']
scores = target['scores']
labels = target['labels']
shape = shape.cpu().numpy()
masks = masks.cpu().numpy()
scores = scores.cpu().numpy()
labels = labels.cpu().numpy()
select_shape = []
select_masks = []
select_scores = []
select_labels = []
for i in range(len(scores)):
if (scores[i]>self.score):
select_shape.append(shape[i])
select_masks.append(masks[i])
select_scores.append(scores[i])
select_labels.append(labels[i])
output = {
'boxes': np.array(select_shape, dtype=np.int32).tolist(),
# 'masks': np.array(select_masks, dtype=np.uint8).tolist(),
'scores': np.array(select_scores, dtype=np.float32).tolist(),
'labels': np.array(select_labels, dtype=np.int32).tolist()
}
self.output.append(output)
def on_predict_end(self):
self.output = self.output |