Spaces:
Paused
Paused
Update CatVTON/model/DensePose/__init__.py
Browse files
CatVTON/model/DensePose/__init__.py
CHANGED
|
@@ -38,22 +38,18 @@ class DensePose:
|
|
| 38 |
self.min_score = 0.8
|
| 39 |
|
| 40 |
self.cfg = self.setup_config()
|
| 41 |
-
self.cfg.MODEL.DEVICE = self.device # Force device before creating predictor
|
| 42 |
self.predictor = DefaultPredictor(self.cfg)
|
| 43 |
|
| 44 |
def setup_config(self):
|
| 45 |
-
opts = [
|
| 46 |
-
"MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score),
|
| 47 |
-
"MODEL.DEVICE", self.device # Add this line
|
| 48 |
-
]
|
| 49 |
cfg = get_cfg()
|
| 50 |
add_densepose_config(cfg)
|
| 51 |
cfg.merge_from_file(self.config_path)
|
| 52 |
cfg.merge_from_list(opts)
|
| 53 |
cfg.MODEL.WEIGHTS = self.model_path
|
|
|
|
| 54 |
cfg.freeze()
|
| 55 |
return cfg
|
| 56 |
-
|
| 57 |
@staticmethod
|
| 58 |
def _get_input_file_list(input_spec: str):
|
| 59 |
if os.path.isdir(input_spec):
|
|
|
|
| 38 |
self.min_score = 0.8
|
| 39 |
|
| 40 |
self.cfg = self.setup_config()
|
|
|
|
| 41 |
self.predictor = DefaultPredictor(self.cfg)
|
| 42 |
|
| 43 |
def setup_config(self):
|
| 44 |
+
opts = ["MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score)]
|
|
|
|
|
|
|
|
|
|
| 45 |
cfg = get_cfg()
|
| 46 |
add_densepose_config(cfg)
|
| 47 |
cfg.merge_from_file(self.config_path)
|
| 48 |
cfg.merge_from_list(opts)
|
| 49 |
cfg.MODEL.WEIGHTS = self.model_path
|
| 50 |
+
cfg.MODEL.DEVICE = self.device # Set device BEFORE freeze
|
| 51 |
cfg.freeze()
|
| 52 |
return cfg
|
|
|
|
| 53 |
@staticmethod
|
| 54 |
def _get_input_file_list(input_spec: str):
|
| 55 |
if os.path.isdir(input_spec):
|