yourusername commited on
Commit
8b2e7f8
1 Parent(s): 573cc9d

:tada: init

Browse files
Files changed (4) hide show
  1. README.md +7 -0
  2. pipeline.py +23 -0
  3. pytorch_model.bin +3 -0
  4. requirements.txt +1 -0
README.md ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - image-classification
4
+ library_name: generic
5
+ ---
6
+
7
+ # Test
pipeline.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ from flash.image import ImageClassifier
4
+
5
+ # ⚠️ You need this to access the state key
6
+ from flash.core.data.data_source import LabelsState
7
+
8
+
9
+ class PreTrainedPipeline():
10
+ def __init__(self, path=""):
11
+ self.device = 'cpu'
12
+ self.model = ImageClassifier.load_from_checkpoint(os.path.join(path, "pytorch_model.bin"))
13
+ self.data_pipeline = self.model.build_data_pipeline()
14
+ self.labels = self.model._data_pipeline_state._state[LabelsState].labels
15
+ self.top_k = 5
16
+
17
+ def __call__(self, inputs):
18
+ x = self.data_pipeline.worker_preprocessor('predict')({'input': inputs})
19
+ x = self.model.transfer_batch_to_device(x, self.device, 0)
20
+ x = self.data_pipeline.device_preprocessor('predict')(x)
21
+ out = self.model.predict_step(x, 0)
22
+ proba = out['preds'].softmax(1)[0].tolist()
23
+ return [{'score': s, 'label': l} for s, l in sorted(zip(proba, self.labels), reverse=True)[:self.top_k]]
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a9b48897cfec1f3038b23cd79630b9f64ed0342e7f50c6d44175b543cb683283
3
+ size 94911265
requirements.txt ADDED
@@ -0,0 +1 @@
 
1
+ lightning-flash[image]